import json
import os

from gxl_ai_utils.utils import utils_file


def main_handle_wenetspeech_get_wavlist():
    raw_dir = "/mnt/apdcephfs_sgfd/share_303841515/Tealab/user/xuelonggeng/data/wenetspeech_data/L_fixed/downloads/extracted"
    output_wav_list_path = "/mnt/apdcephfs_sgfd/share_303841515/Tealab/user/xuelonggeng/data/wenetspeech_data/wav.list"
    output_wav_scp_path = "/mnt/apdcephfs_sgfd/share_303841515/Tealab/user/xuelonggeng/data/wenetspeech_data/wav.scp"

    # 1. 获取一级目录下的所有直接子目录（排除文件）
    first_level_dirs = []
    for entry in os.scandir(raw_dir):
        if entry.is_dir(follow_symlinks=False):  # 不跟随符号链接，避免循环引用
            first_level_dirs.append(entry.path)

    print(f"一级目录共 {len(first_level_dirs)} 个：")


    # 2. 对于每个一级目录，获取其下所有"叶子目录"（无更深子目录的目录）
    leaf_dir_list = []
    for first_dir in utils_file.tqdm(first_level_dirs, desc="处理一级目录,get leaf dirs", total=len(first_level_dirs)):
        leaf_dirs = []
        # 递归遍历当前一级目录下的所有子目录
        for root, dirs, files in os.walk(first_dir):
            # 若当前目录没有子目录（dirs为空），则为叶子目录
            if not dirs:
                leaf_dirs.append(root)
        leaf_dir_list.extend(leaf_dirs)

    print(f"叶子目录共 {len(leaf_dir_list)} 个：")
    all_list = []
    for leaf_dir_path  in utils_file.tqdm(leaf_dir_list, desc="处理叶子目录,get leaf dirs", total=len(leaf_dir_list)):
        wav_list = utils_file.do_get_list_for_wav_dir(leaf_dir_path)
        all_list.extend(wav_list)
    utils_file.write_list_to_file(all_list, output_wav_list_path)
    wav_dict = {}
    for wav_path in all_list:
        key = utils_file.do_get_file_pure_name_from_path(wav_path)
        wav_dict[key] = wav_path
    utils_file.write_dict_to_scp(wav_dict, output_wav_scp_path)

def main_wenetspeech_get_text():
    json_path = "/mnt/apdcephfs_sgfd/share_303841515/Tealab/user/xuelonggeng/data/wenetspeech_data/ft_local/WenetSpeech.json"
    json_dict = json.load(open(json_path, 'r', encoding='utf-8'))
    dataset = json_dict["dataset"]
    print(f'load successfully {len(dataset)}')
    big_res_dict = {}
    audios_list = json_dict["audios"]
    for audio_item_dict in utils_file.tqdm(audios_list, desc="处理音频列表", total=len(audios_list)):
        segments_list = audio_item_dict["segments"]
        for segment_item_dict in segments_list:
            sid = segment_item_dict["sid"]
            subsets = segment_item_dict["subsets"]
            text = segment_item_dict["text"]
            for subset_item in subsets:
                if subset_item not in big_res_dict:
                    big_res_dict[subset_item] = {}
                big_res_dict[subset_item][sid] = text
    output_text_path = "/mnt/apdcephfs_sgfd/share_303841515/Tealab/user/xuelonggeng/data/wenetspeech_data/text_{}.scp"
    for subset_name, subset_dict in big_res_dict.items():
        utils_file.write_dict_to_scp(subset_dict, output_text_path.format(subset_name))


def main_get_all_datalist():
    """"""
    root_dir = "/mnt/apdcephfs_sgfd/share_303841515/Tealab/user/xuelonggeng/data/wenetspeech_data"
    text_S_path = f'{root_dir}/text_S.scp'
    text_M_path = f'{root_dir}/text_M.scp'
    text_L_path = f'{root_dir}/text_L.scp'
    text_DEV_path = f'{root_dir}/text_DEV.scp'
    text_W_path = f'{root_dir}/text_W.scp'
    text_text_net_path = f'{root_dir}/text_TEST_NET.scp'
    text_test_meeting_path = f'{root_dir}/text_TEST_MEETING.scp'
    wav_scp_path = f'{root_dir}/wav.scp'
    wav_dict = utils_file.load_dict_from_scp(wav_scp_path)
    text_S_dict = utils_file.load_dict_from_scp(text_S_path)
    text_M_dict = utils_file.load_dict_from_scp(text_M_path)
    text_L_dict = utils_file.load_dict_from_scp(text_L_path)
    text_DEV_dict = utils_file.load_dict_from_scp(text_DEV_path)
    text_W_dict = utils_file.load_dict_from_scp(text_W_path)
    text_text_net_dict = utils_file.load_dict_from_scp(text_text_net_path)
    text_test_meeting_dict = utils_file.load_dict_from_scp(text_test_meeting_path)
    wav_S_dict = {}
    wav_M_dict = {}
    wav_L_dict = {}
    wav_DEV_dict = {}
    wav_W_dict = {}
    wav_text_net_dict = {}
    wav_test_meeting_dict = {}
    for key, value in utils_file.tqdm(wav_dict.items(), desc="处理音频列表", total=len(wav_dict)):
        if key in text_S_dict:
            wav_S_dict[key] = value
        if key in text_M_dict:
            wav_M_dict[key] = value
        if key in text_L_dict:
            wav_L_dict[key] = value
        if key in text_DEV_dict:
            wav_DEV_dict[key] = value
        if key in text_W_dict:
            wav_W_dict[key] = value
        if key in text_text_net_dict:
            wav_text_net_dict[key] = value
        if key in text_test_meeting_dict:
            wav_test_meeting_dict[key] = value
    dict_list_S = []
    dict_list_M = []
    dict_list_L = []
    dict_list_DEV = []
    dict_list_W = []
    dict_list_text_net = []
    dict_list_test_meeting = []
    for key, value in utils_file.tqdm(wav_S_dict.items(), desc="处理音频列表", total=len(wav_S_dict)):
        dict_list_S.append({'key': key, "wav": value, "text": text_S_dict[key]})
    for key, value in utils_file.tqdm(wav_M_dict.items(), desc="处理音频列表", total=len(wav_M_dict)):
        dict_list_M.append({'key': key, "wav": value, "text": text_M_dict[key]})
    for key, value in utils_file.tqdm(wav_L_dict.items(), desc="处理音频列表", total=len(wav_L_dict)):
        dict_list_L.append({'key': key, "wav": value, "text": text_L_dict[key]})
    for key, value in utils_file.tqdm(wav_DEV_dict.items(), desc="处理音频列表", total=len(wav_DEV_dict)):
        dict_list_DEV.append({'key': key, "wav": value, "text": text_DEV_dict[key]})
    for key, value in utils_file.tqdm(wav_W_dict.items(), desc="处理音频列表", total=len(wav_W_dict)):
        dict_list_W.append({'key': key, "wav": value, "text": text_W_dict[key]})
    for key, value in utils_file.tqdm(wav_text_net_dict.items(), desc="处理音频列表", total=len(wav_text_net_dict)):
        dict_list_text_net.append({'key': key, "wav": value, "text": text_text_net_dict[key]})
    for key, value in utils_file.tqdm(wav_test_meeting_dict.items(), desc="处理音频列表", total=len(wav_test_meeting_dict)):
        dict_list_test_meeting.append({'key': key, "wav": value, "text": text_test_meeting_dict[key]})
    wav_S_path = f'{root_dir}/wav_S.scp'
    wav_M_path = f'{root_dir}/wav_M.scp'
    wav_L_path = f'{root_dir}/wav_L.scp'
    wav_DEV_path = f'{root_dir}/wav_DEV.scp'
    wav_W_path = f'{root_dir}/wav_W.scp'
    wav_text_net_path = f'{root_dir}/wav_TEST_NET.scp'
    wav_test_meeting_path = f'{root_dir}/wav_TEST_MEETING.scp'
    utils_file.write_dict_to_scp(wav_S_dict, wav_S_path)
    utils_file.write_dict_to_scp(wav_M_dict, wav_M_path)
    utils_file.write_dict_to_scp(wav_L_dict, wav_L_path)
    utils_file.write_dict_to_scp(wav_DEV_dict, wav_DEV_path)
    utils_file.write_dict_to_scp(wav_W_dict, wav_W_path)
    utils_file.write_dict_to_scp(wav_text_net_dict, wav_text_net_path)
    utils_file.write_dict_to_scp(wav_test_meeting_dict, wav_test_meeting_path)
    dict_list_S_path = f'{root_dir}/data_S.jsonl'
    dict_list_M_path = f'{root_dir}/data_M.jsonl'
    dict_list_L_path = f'{root_dir}/data_L.jsonl'
    dict_list_DEV_path = f'{root_dir}/data_DEV.jsonl'
    dict_list_W_path = f'{root_dir}/data_W.jsonl'
    dict_list_text_net_path = f'{root_dir}/data_TEST_NET.jsonl'
    dict_list_test_meeting_path = f'{root_dir}/data_TEST_MEETING.jsonl'
    utils_file.write_dict_list_to_jsonl(dict_list_S, dict_list_S_path)
    utils_file.write_dict_list_to_jsonl(dict_list_M, dict_list_M_path)
    utils_file.write_dict_list_to_jsonl(dict_list_L, dict_list_L_path)
    utils_file.write_dict_list_to_jsonl(dict_list_DEV, dict_list_DEV_path)
    utils_file.write_dict_list_to_jsonl(dict_list_W, dict_list_W_path)
    utils_file.write_dict_list_to_jsonl(dict_list_text_net, dict_list_text_net_path)
    utils_file.write_dict_list_to_jsonl(dict_list_test_meeting, dict_list_test_meeting_path)
if __name__ == "__main__":
    my_timer = utils_file.GxlTimer()
    my_timer.start()
    main_get_all_datalist()
    my_timer.stop_halfway_and_print()